package com.zrj.gateway.filter;

import com.alibaba.fastjson2.JSON;
import lombok.extern.slf4j.Slf4j;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.cloud.gateway.support.BodyInserterContext;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.http.HttpHeaders;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.stereotype.Component;
import org.springframework.web.reactive.function.BodyInserter;
import org.springframework.web.reactive.function.BodyInserters;
import org.springframework.web.reactive.function.server.HandlerStrategies;
import org.springframework.web.reactive.function.server.ServerRequest;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.util.Map;
import java.util.function.BiFunction;

/**
 * 请求过滤器
 * 参考 {@link org.springframework.cloud.gateway.filter.factory.rewrite.ModifyRequestBodyGatewayFilterFactory}
 *
 * @author zrj
 * @since 2022/10/7
 **/
@Slf4j
@Component
public class ModifyRequestGlobalFilter implements GlobalFilter, Ordered {

    @Override
    public int getOrder() {
        return -2000;
    }

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {

        ServerRequest serverRequest = ServerRequest.create( exchange, HandlerStrategies.withDefaults().messageReaders() );
        Mono<String> modifiedBody = serverRequest.bodyToMono( String.class ).flatMap( originalBody -> modifyBody()
                .apply( exchange, originalBody ) );

        BodyInserter bodyInserter = BodyInserters.fromPublisher( modifiedBody, String.class );

        HttpHeaders headers = new HttpHeaders();
        headers.putAll( exchange.getRequest().getHeaders() );
        headers.remove( "Content-Length" );

        ModifyCachedBodyOutputMessage outputMessage = new ModifyCachedBodyOutputMessage( exchange, headers );

        return bodyInserter.insert( outputMessage, new BodyInserterContext() ).then( Mono.defer( () -> {
            ServerHttpRequest decorator = decorate( exchange, headers, outputMessage );
            return chain.filter( exchange.mutate().request( decorator ).build() );
        } ) ).onErrorResume( (throwable) -> release( exchange, outputMessage, (Throwable) throwable ) );
    }

    private BiFunction<ServerWebExchange, String, Mono<String>> modifyBody() {
        return new BiFunction<ServerWebExchange, String, Mono<String>>() {
            @Override
            public Mono<String> apply(ServerWebExchange serverWebExchange, String raw) {
                try {
                    //String decrypt = AesUtils.decrypt( raw );
                    //log.info( "修改请求体,修改前:{},修改后:{}", raw, decrypt );
                    //return Mono.just( decrypt );

                    Map<String, Object> requestMap = JSON.parseObject( raw, Map.class );
                    log.info( "【请求过滤器】修改请求体,修改前:{},修改后:{}", requestMap );
                    return Mono.just( raw );
                } catch (Exception e) {
                    log.error( "请求过滤器服务器异常", e );
                    return Mono.empty();
                }
            }
        };
    }


    protected Mono<Void> release(ServerWebExchange exchange, ModifyCachedBodyOutputMessage outputMessage, Throwable throwable) {
        return outputMessage.isCached() ? outputMessage.getBody().map( DataBufferUtils::release ).then( Mono.error( throwable ) ) : Mono.error( throwable );
    }

    ServerHttpRequestDecorator decorate(ServerWebExchange exchange, HttpHeaders headers, ModifyCachedBodyOutputMessage outputMessage) {
        return new ServerHttpRequestDecorator( exchange.getRequest() ) {
            @Override
            public HttpHeaders getHeaders() {
                long contentLength = headers.getContentLength();
                HttpHeaders httpHeaders = new HttpHeaders();
                httpHeaders.putAll( headers );
                if (contentLength > 0L) {
                    httpHeaders.setContentLength( contentLength );
                } else {
                    httpHeaders.set( "Transfer-Encoding", "chunked" );
                }

                return httpHeaders;
            }

            @Override
            public Flux<DataBuffer> getBody() {
                return outputMessage.getBody();
            }
        };
    }
}

